/*
 * Copyright (c) 2022 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "hks_access_control_test_common.h"

#include <gtest/gtest.h>

using namespace testing::ext;
namespace Unittest::HksAccessControlPartTest {
static struct HksParam AuthToken_Import_Params[] = {
    { 
        .tag = HKS_TAG_ALGORITHM, 
        .uint32Param = HKS_ALG_HMAC 
    },{ 
        .tag = HKS_TAG_PURPOSE, 
        .uint32Param = HKS_KEY_PURPOSE_MAC 
    },{ 
        .tag = HKS_TAG_KEY_SIZE, 
        .uint32Param = HKS_AES_KEY_SIZE_256 
    },{ 
        .tag = HKS_TAG_DIGEST, 
        .uint32Param = HKS_DIGEST_SHA256 
    }
};

static struct HksParam AuthToken_HMAC_Params[] = {
    {
        .tag = HKS_TAG_ALGORITHM,
        .uint32Param = HKS_ALG_HMAC
    }, {
        .tag = HKS_TAG_PURPOSE,
        .uint32Param = HKS_KEY_PURPOSE_MAC
    }, {
        .tag = HKS_TAG_DIGEST,
        .uint32Param = HKS_DIGEST_SHA256
    }
};

int32_t AuthTokenImportKey()
{
    struct HksParamSet *importParamSet = nullptr;
    int32_t ret = InitParamSet(&importParamSet, AuthToken_Import_Params, sizeof(AuthToken_Import_Params) / sizeof(HksParam));
    if (ret != HKS_SUCCESS) {
        return ret;
    }
    struct HksBlob key = { SHA256_KEY_LEN, (uint8_t*)HKS_DEFAULT_USER_AT_KEY };

    uint8_t alias[] = "AuthToken_Sign_Verify_KeyAlias";
    struct HksBlob keyAlias = { sizeof(alias), alias };
    ret = HksImportKey(&keyAlias, importParamSet, &key);
    HksFreeParamSet(&importParamSet);
    return ret;
}

int32_t AuthTokenSign(struct HksBlob *challenge, uint64_t secureUid, uint64_t enrolledId, uint64_t time, std::vector<uint8_t>& token)
{
    int32_t ret = HKS_FAILURE;
    uint8_t alias[] = "AuthToken_Sign_Verify_KeyAlias";
    const struct HksBlob keyAlias = { sizeof(alias), alias };
    AuthTokenImportKey();
    HksUserAuthToken *authTokenHal = (struct HksUserAuthToken *)HksMalloc(AUTH_TOKEN_LEN);
    if (authTokenHal == NULL) {
        HKS_LOG_E("could not alloc memory");
        return HKS_FAILURE;
    }
    uint64_t curTime = 0;
            HKS_LOG_E("AuthTokenSign 0");
    ret = HksCoreHalElapsedRealTime(&curTime);
    if (ret != HKS_SUCCESS) {
        HKS_LOG_E("get elapsed real time failed!");
        return ret;
    }
    authTokenHal->time = curTime + time;
    authTokenHal->secureUid = secureUid;
    authTokenHal->enrolledId = enrolledId;
        HKS_LOG_E("AuthTokenSign 1");

    if (memcpy_s(authTokenHal->challenge, TOKEN_CHALLENGE_LEN, challenge->data, challenge->size) != EOK) {
        return HKS_FAILURE;
    }
    uint8_t authTokenData[AUTH_TOKEN_DATA_LEN] = {0};
    if (memcpy_s(authTokenData, AUTH_TOKEN_DATA_LEN, authTokenHal, AUTH_TOKEN_DATA_LEN) != EOK) {
        return HKS_FAILURE;
    }
    struct HksBlob inData = { AUTH_TOKEN_DATA_LEN, authTokenData };

    HKS_LOG_E("AuthTokenSign 2");

    struct HksParamSet *hmacParamSet = nullptr;
    ret = InitParamSet(&hmacParamSet, AuthToken_HMAC_Params, sizeof(AuthToken_HMAC_Params) / sizeof(HksParam));
    if (ret != HKS_SUCCESS) {
        HKS_LOG_E("InitsignParamSet failed");
        return ret;
    }
    /// Init
    uint8_t handle[32] = {0};
    struct HksBlob handleHMAC = { 32, handle };

        HKS_LOG_E("AuthTokenSign 3");

    ret = HksInit(&keyAlias, hmacParamSet, &handleHMAC, nullptr);
    if (ret != HKS_SUCCESS) {
        HKS_LOG_E("Init failed");
        HksFreeParamSet(&hmacParamSet);
        return ret;
    }
    // Update & Finish
    struct HksBlob outData = { SHA256_SIGN_LEN, authTokenHal->sign };

        HKS_LOG_E("AuthTokenSign 4");

    ret = TestUpdateFinish(&handleHMAC, hmacParamSet, HKS_KEY_PURPOSE_MAC, &inData, &outData);
    if (ret != HKS_SUCCESS) {
        HKS_LOG_E("Update&Finish failed");
        HksFreeParamSet(&hmacParamSet);
        return ret;
    }
        HKS_LOG_E("AuthTokenSign 5");

    for (uint32_t i = 0; i < AUTH_TOKEN_DATA_LEN; i++){
        token.push_back(authTokenData[i]);
    }
    for (uint32_t i = 0; i < SHA256_SIGN_LEN; i++){
        token.push_back(authTokenHal->sign[i]);
    }
    HksFreeParamSet(&hmacParamSet);
    return  HKS_SUCCESS;
}

static int32_t AppendToNewParamSet(const struct HksParamSet *paramSet, struct HksParamSet **outParamSet)
{
    int32_t ret;
    struct HksParamSet *newParamSet = NULL;
    HKS_LOG_E("AppendToNewParamSet start ");
    do {
        ret = HksCheckParamSet(paramSet, paramSet->paramSetSize);
        if (ret != HKS_SUCCESS) {
            HKS_LOG_E("check paramSet failed");
            break;
        }

        ret = HksFreshParamSet((struct HksParamSet *)paramSet, false);
        if (ret != HKS_SUCCESS) {
            HKS_LOG_E("append fresh paramset failed");
            break;
        }

        ret = HksInitParamSet(&newParamSet);
        if (ret != HKS_SUCCESS) {
            HKS_LOG_E("append init operation param set failed");
            break;
        }

        ret = HksAddParams(newParamSet, paramSet->params, paramSet->paramsCnt);
        if (ret != HKS_SUCCESS) {
            HKS_LOG_E("append params failed");
            break;
        }
    HKS_LOG_E("AppendToNewParamSet end ");

        *outParamSet = newParamSet;
        return ret;
    } while (0);

    HksFreeParamSet(&newParamSet);
    return ret;
}

 int32_t HksBuildAuthtoken(struct HksParamSet **initParamSet, struct HksBlob *authChallenge, uint64_t secureUid, 
    uint64_t enrolledId, uint64_t time){

    struct HksParam tmpParams;
    std::vector<uint8_t> token;
    int32_t ret = AuthTokenSign(authChallenge, secureUid, enrolledId, time, token);
    if (ret != HKS_SUCCESS) {
        HKS_LOG_E("AuthToeknSign Failed.");
        return ret;
    }
    HKS_LOG_I("HksBuildAuthtoken 1.");
    uint8_t authToken[AUTH_TOKEN_LEN] = {0};
    for (int i = 0; i < AUTH_TOKEN_LEN; i++) {
        authToken[i] = token[i]; 
    }
    tmpParams.tag = HKS_TAG_AUTH_TOKEN;
    tmpParams.blob.data = authToken;
    tmpParams.blob.size = AUTH_TOKEN_LEN;
    HKS_LOG_I("AuthToekn Data: %u", authToken);

    struct HksParamSet *newParamSet = NULL;

    ret = AppendToNewParamSet(*initParamSet, &newParamSet);
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("AppendToNewParamSet failed!\n");
        return ret;
    }

    HKS_LOG_I("HksBuildAuthtoken 1.5");

    ret = HksAddParams(newParamSet, &tmpParams, 1);
    if (ret != 0) {
        HKS_LOG_I("HksAddParam failed!\n");
        return ret;
    }
    ret = HksBuildParamSet(&newParamSet);
    if (ret != 0) {
        HKS_LOG_I("HksBuildParamSet failed!\n");
        return ret;
    }
    
    *initParamSet = newParamSet;

    HKS_LOG_I("HksBuildAuthtoken 3.");

    return ret;
}

int32_t AddAuthtokenUpdateFinish(struct HksBlob *handle, 
    struct HksParamSet *initParamSet, uint32_t posNum){
    struct HksBlob inData = { g_inData.length(), (uint8_t *)g_inData.c_str() };
    uint8_t outDataS[DATA_COMMON_SIZE] = {0};
    struct HksBlob outDataSign = { DATA_COMMON_SIZE, outDataS };

    struct HksParam challengePos;
    challengePos.tag = HKS_TAG_CHALLENGE_POS;
    challengePos.uint32Param = posNum;
    
    int ret = HksAddParams(initParamSet, &challengePos, 1);

    ret = TestUpdateLoopFinish(handle, initParamSet, &inData, &outDataSign);
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("TestUpdateLoopFinish failed, ret : %d", ret);
    }
    return ret;
}

int32_t CheckAccessCipherTest(const TestAccessCaseParams &testCaseParams, 
uint64_t secureUid, uint64_t enrolledId, uint64_t time)
{
    struct HksParamSet *genParamSet = nullptr;

    int32_t ret = InitParamSet(&genParamSet, testCaseParams.genParams.data(), testCaseParams.genParams.size());
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("InitParamSet(gen) failed, ret : %d", ret);
        return ret;
    }
    uint8_t alias[] = "testCheckAuth";
    struct HksBlob keyAlias = { sizeof(alias), alias };
    ret = HksGenerateKey(&keyAlias, genParamSet, nullptr);
    if (ret != HKS_SUCCESS) {
        HksFreeParamSet(&genParamSet);
        HKS_LOG_I("HksGenerateKey failed, ret : %d", ret);
        return ret;
    }

    struct HksParamSet *initParamSet = nullptr;
    ret = InitParamSet(&initParamSet, testCaseParams.initParams.data(), testCaseParams.initParams.size());
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("InitParamSet(init) failed, ret : %d", ret);
        return ret;
    }
    uint8_t challenge[sizeof(uint64_t)] = {0}; 
    struct HksBlob challengeBlob = { sizeof(uint64_t), challenge };

    uint8_t tmpHandle[sizeof(uint64_t)] = {0};
    struct HksBlob handle = { sizeof(uint64_t), tmpHandle };
	ret = HksInit(&keyAlias, initParamSet, &handle, &challengeBlob);
    if (ret != HKS_SUCCESS) {
        return ret;
    }
    
    uint8_t authChallenge[32] = {0};
    (void) memcpy_s(&initParamSet, sizeof(uint64_t), authChallenge, sizeof(uint64_t));
    ret = HksBuildAuthtoken(&initParamSet, &challengeBlob, secureUid, enrolledId, time);

    ret = AddAuthtokenUpdateFinish(&handle, initParamSet, 0);
    
    HksFreeParamSet(&genParamSet);
    HksFreeParamSet(&initParamSet);
    (void)HksDeleteKey(&keyAlias, nullptr);

    return (ret == testCaseParams.initResult) ? HKS_SUCCESS : HKS_FAILURE;
}

int32_t CheckAccessHmacTest(const TestAccessCaseParams &testCaseParams, 
uint64_t secureUid, uint64_t enrolledId, uint64_t time)
{
    struct HksParamSet *genParamSet = nullptr;
    int32_t ret = InitParamSet(&genParamSet, testCaseParams.genParams.data(), testCaseParams.genParams.size());
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("InitParamSet(gen) failed, ret : %d", ret);
        return ret;
    }
    uint8_t alias[] = "testCheckAuth";
    struct HksBlob keyAlias = { sizeof(alias), alias };
    ret = HksGenerateKey(&keyAlias, genParamSet, nullptr);
    if (ret != HKS_SUCCESS) {
        HksFreeParamSet(&genParamSet);
        HKS_LOG_I("HksGenerateKey failed, ret : %d", ret);
        return ret;
    }

    struct HksParamSet *initParamSet = nullptr;
    ret = InitParamSet(&initParamSet, testCaseParams.initParams.data(), testCaseParams.initParams.size());
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("InitParamSet(init) failed, ret : %d", ret);
        return ret;
    }
    uint8_t challenge[sizeof(uint64_t)] = {0}; 
    struct HksBlob challengeBlob = { sizeof(uint64_t), challenge };
    
    uint8_t handle[sizeof(uint64_t)] = {0};
    struct HksBlob handleHMAC = { sizeof(uint64_t), handle };
    ret = HksInit(&keyAlias, initParamSet, &handleHMAC, &challengeBlob);
    if (ret != HKS_SUCCESS) {
        HksDeleteKey(&keyAlias, genParamSet);
        return ret;
    }

    uint8_t authChallenge[32] = {0};
    (void) memcpy_s(&initParamSet, sizeof(uint64_t), authChallenge, sizeof(uint64_t));
    ret = HksBuildAuthtoken(&initParamSet, &challengeBlob, secureUid, enrolledId, time);
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("HksBuildAuthtoken failed, ret : %d", ret);
        HksDeleteKey(&keyAlias, genParamSet);
        return ret;
    }
    // Update & Finish
    struct HksBlob inData = { g_inData.length(), (uint8_t *)g_inData.c_str() };
    uint8_t out[HMAC_COMMON_SIZE] = {0};
    struct HksBlob outData = { HMAC_COMMON_SIZE, out };
    ret = TestUpdateFinish(&handleHMAC, initParamSet, HKS_KEY_PURPOSE_MAC, &inData, &outData);
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("TestUpdateFinish failed, ret : %d", ret);
        HksDeleteKey(&keyAlias, genParamSet);
        return ret;
    }

    /* 3. Delete Key */
    ret = HksDeleteKey(&keyAlias, genParamSet);
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("HksDeleteKey failed, ret : %d", ret);
        HksDeleteKey(&keyAlias, genParamSet);
        return ret;
    }

    return (ret == testCaseParams.initResult) ? HKS_SUCCESS : HKS_FAILURE;
}

int32_t CheckAccessAgreeTest(const TestAccessCaseParams &testCaseParams, const struct HksParamSet *finishParamSet, 
uint64_t secureUid, uint64_t enrolledId, uint64_t time)
{
    struct HksParamSet *genParamSet = nullptr;
    int32_t ret = InitParamSet(&genParamSet, testCaseParams.genParams.data(), testCaseParams.genParams.size());
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("InitParamSet(gen) failed, ret : %d", ret);
        return ret;
    }
    uint8_t alias[] = "testCheckAuth";
    struct HksBlob keyAlias = { sizeof(alias), alias };
    ret = HksGenerateKey(&keyAlias, genParamSet, nullptr);
    if (ret != HKS_SUCCESS) {
        HksFreeParamSet(&genParamSet);
        HKS_LOG_I("HksGenerateKey failed, ret : %d", ret);
        return ret;
    }

    uint8_t pubKey[4096] = {0}; 
    struct HksBlob publicKey = { 4096, pubKey };
    ret = HksExportPublicKey(&keyAlias, genParamSet, &publicKey);
    EXPECT_EQ(ret, HKS_SUCCESS) << "ExportPublicKey01 failed.";
    if (ret != HKS_SUCCESS) {
        return HKS_FAILURE;
    }

    struct HksBlob inData = { g_inData.length(), (uint8_t *)g_inData.c_str() };
    struct HksParamSet *initParamSet = nullptr;
    ret = InitParamSet(&initParamSet, testCaseParams.initParams.data(), testCaseParams.initParams.size());
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("InitParamSet(init) failed, ret : %d", ret);
        return ret;
    }
    uint8_t challenge[32] = {0}; 
    struct HksBlob challengeBlob = { 32, challenge };

    uint8_t handleU[32] = {0};
    struct HksBlob handle = { 32, handleU };
    ret = HksInit(&keyAlias, initParamSet, &handle, &challengeBlob);
    if (ret != HKS_SUCCESS) {
        return HKS_FAILURE;
    }

    //uint8_t authChallenge[32] = {0};
    //(void) memcpy_s(&initParamSet, sizeof(uint64_t), authChallenge, sizeof(uint64_t));

    ret = HksBuildAuthtoken(&initParamSet, &challengeBlob, secureUid, enrolledId, time);
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("HksBuildAuthtoken failed, ret : %d", ret);
        HksDeleteKey(&keyAlias, genParamSet);
        return ret;
    }

    uint8_t outDataU[ECDH_COMMON_SIZE] = {0};
    struct HksBlob outDataUpdate = { ECDH_COMMON_SIZE, outDataU };
    ret = HksUpdate((const struct HksBlob *)&handle, initParamSet, &publicKey, &outDataUpdate);
    EXPECT_EQ(ret, HKS_SUCCESS) << "Update failed.";
    if (ret != HKS_SUCCESS) {
        return HKS_FAILURE;
    }

    uint8_t outDataF[ECDH_COMMON_SIZE] = {0};
    struct HksBlob outDataFinish = { ECDH_COMMON_SIZE, outDataF };
    ret = HksFinish(&handle, finishParamSet, (const struct HksBlob *)&inData, &outDataFinish);
    EXPECT_EQ(ret, HKS_SUCCESS) << "Finish failed.";
    if (ret != HKS_SUCCESS) {
        return HKS_FAILURE;
    }
    
    return (ret == testCaseParams.initResult) ? HKS_SUCCESS : HKS_FAILURE;
}

int32_t CheckAccessDeriveTest(const TestAccessCaseParams &testCaseParams, const struct HksParamSet *finishParamSet, 
uint64_t secureUid, uint64_t enrolledId, uint64_t time)
{
    struct HksBlob inData = { g_inData.length(), (uint8_t *)g_inData.c_str() };
    struct HksParamSet *genParamSet = nullptr;
    struct HksParamSet *initParamSet = nullptr;
    int32_t ret = InitParamSet(&genParamSet, testCaseParams.genParams.data(), testCaseParams.genParams.size());
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("InitParamSet(gen) failed, ret : %d", ret);
        return ret;
    }
    uint8_t alias[] = "testCheckAuth";
    struct HksBlob keyAlias = { sizeof(alias), alias };
    ret = HksGenerateKey(&keyAlias, genParamSet, nullptr);
    if (ret != HKS_SUCCESS) {
        HksFreeParamSet(&genParamSet);
        HKS_LOG_I("HksGenerateKey failed, ret : %d", ret);
        return ret;
    }

    ret = InitParamSet(&initParamSet, testCaseParams.initParams.data(), testCaseParams.initParams.size());
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("InitParamSet(init) failed, ret : %d", ret);
        return ret;
    }
    uint8_t challenge[sizeof(uint64_t)] = {0}; 
    struct HksBlob challengeBlob = { sizeof(uint64_t), challenge };
    // Init
    uint8_t handleD[sizeof(uint64_t)] = {0};
    struct HksBlob handleDerive = { sizeof(uint64_t), handleD };
    ret = HksInit(&keyAlias, initParamSet, &handleDerive, &challengeBlob);
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("Init failed., ret : %d", ret);
        HksDeleteKey(&keyAlias, genParamSet);
        return ret;
    }
    uint8_t authChallenge[32] = {0};
    (void) memcpy_s(&initParamSet, sizeof(uint64_t), authChallenge, sizeof(uint64_t));
    ret = HksBuildAuthtoken(&initParamSet, &challengeBlob, secureUid, enrolledId, time);
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("HksBuildAuthtoken failed, ret : %d", ret);
        HksDeleteKey(&keyAlias, genParamSet);
        return ret;
    }
    // Update
    uint8_t tmpOut[DERIVE_COMMON_SIZE] = {0};
    struct HksBlob outData = { DERIVE_COMMON_SIZE, tmpOut };
    ret = HksUpdate(&handleDerive, initParamSet, &inData, &outData);
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("Update failed., ret : %d", ret);
        HksDeleteKey(&keyAlias, genParamSet);
        return ret;
    }
    // Finish
    uint8_t outDataD[DERIVE_COMMON_SIZE] = {0};
    struct HksBlob outDataDerive = { DERIVE_COMMON_SIZE, outDataD };
    ret = HksFinish(&handleDerive, finishParamSet, &inData, &outDataDerive);
    if (ret != HKS_SUCCESS) {
        HKS_LOG_I("Finish failed., ret : %d", ret);
        HksDeleteKey(&keyAlias, genParamSet);
        return ret;
    }
    
    HksDeleteKey(&keyAlias, genParamSet);
    return (ret == testCaseParams.initResult) ? HKS_SUCCESS : HKS_FAILURE;
}
}
